package telex;
/*
 * Gale1TNGramDatabase.java
 *
 * Copyright (c) 2007-2008 Colin Bayer, Douglas Downey, Oren Etzioni,
 *   University of Washington Computer Science and Engineering
 * 
 * See the file "COPYING" in the root directory of this distribution for
 * full licensing information.
 */

import java.io.*;
import java.net.URL;
import java.net.URLConnection;
import java.util.*;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.zip.*;


public final class Gale1TNGramDatabase extends NGramDatabase {

    private static boolean writeCounts = false;

    public Gale1TNGramDatabase(String aPath, int aMaxArity) throws FileNotFoundException,
            IOException {
        mPath = aPath;
        mMaxArity = aMaxArity;

        // generate indices to n-gram files.
        mIndices = new QueryString[mMaxArity + 1][];
        for (int i = 1; i <= mMaxArity; i++) {
            if (i != 1) {
                generateIndex(i);
            }
        }

        // reset file pointers for each arity.
        mFilePointers = new Pointer[mMaxArity + 1];
        for (int i = 0; i < mMaxArity + 1; i++) {
            mFilePointers[i] = new Pointer();
        }
        
        // open the 1-gram file right now.
        if (!openFile(1, 0)) {
            throw new IOException();
        }
        
        // load or generate counts now, too...
        if (!loadCounts()) {
            throw new IOException();
        }
    }

    public String generatePath(int aN, int aIdx) {
        if (aIdx < 0) {
            // generate path to n-gram index.
            return mPath + File.separator + aN + "gms" + File.separator + aN + "gm.idx";
        } else if (aN == 1) {
            return mPath + File.separator + "1gms" + File.separator + "vocab.gz";
        } else {
            // generate path to aIdx-th n-gram file.
            return mPath + File.separator + aN + "gms" + File.separator + aN + "gm-"
                    + String.format("%1$04d", aIdx) + ".gz";
        }
    }

    private boolean openFile(int aN, int aIdx) {
        if (aN == 1 && aIdx != 0) {
            return false;
        } else if (aN != 1 && (aIdx < 0 || aIdx > mIndices[aN].length)) {
            return false;
        } else if (aN != 1 && aIdx == mFilePointers[aN].fidx) {
            return true;
        }
        try {
            if (mFilePointers[aN].fp != null) {
                mFilePointers[aN].fp.close();
            }
            BufferedReader r = new BufferedReader(new InputStreamReader(new GZIPInputStream(new FileInputStream(generatePath(aN, aIdx)))));

            mFilePointers[aN].fp = r;
            mFilePointers[aN].fidx = aIdx;
        } catch (Exception e) {
            System.out.println("Opening file " + aIdx + " of " + aN + "-grams failed (" + e + ")");
            return false;
        }
        return true;
    }

    private boolean generateIndex(int aN) throws FileNotFoundException, IOException {
        // assumption: the index is sorted in order of file number, and there
        // are no gaps.
        String path = generatePath(aN, -1);

        BufferedReader r;
        ArrayList<QueryString> first_ngs = new ArrayList<QueryString>();
        r = new BufferedReader(new FileReader(path));


        String l;
        int i = 0;

        while ((l = r.readLine()) != null) {
            //String[] w = l.split("\t");
            String[] w = new String[2];
            w[0] = l.substring(0, l.indexOf("\t"));
            w[1] = l.substring(l.indexOf("\t") + 1);
            if (w.length != 2) {
                System.err.println("Warning: " + i + "th entry in " + aN
                        + "-gram index has incorrect format.");
            }
            //System.out.println("dbg: read " + i + "th " + aN +
            //	"-gram entry: \"" + w[0] + "\" / \"" + w[1] + "\"");
            //first_ngs.add(new QueryString(w[1].split(" ")));
            first_ngs.add(new QueryString(w[1].split("\t")));
            i++;
        }

        mIndices[aN] = first_ngs.toArray(new QueryString[0]);
        return true;
    }

    private boolean loadCounts() {
        File f = new File(mPath, "counts.txt");


        mCounts = new long[mMaxArity + 1][];

        try {
            boolean bad_counts = false;

            if (f.exists()) {
                // load counts from file.
                FileReader fr = new FileReader(f);
                BufferedReader r = new BufferedReader(fr);
                
                for (int i = 1; i <= mMaxArity; i++) {
                    mCounts[i] = new long[2];

                    String l = r.readLine();

                    if (l == null) {
                        bad_counts = true;
                        break;
                    }

                    String[] w = l.split("\t");

                    mCounts[i][0] = Long.parseLong(w[0]);
                    mCounts[i][1] = Long.parseLong(w[1]);
                }
                fr.close();
                r.close();
            } else {
                bad_counts = true;
            }

            if (bad_counts) {
                System.out.println("Counts file nonexistent or malformed; beginning a (very slow) count pass.  Please wait patiently (may take several hours).");
                //adicionei linha para ele escrever os counts
                writeCounts = true;
                // compute counts, write counts to file, store in mCounts.
                BufferedWriter w = null;
                if (writeCounts) {
                       w = new BufferedWriter(new FileWriter(f));
                }
                long[] one_gram_counts = null;
                for (int i = 1; i <= mMaxArity; i++) {
                    mCounts[i] = fastReadCount(i, one_gram_counts);
                    if (i == 1) {
                        one_gram_counts = mCounts[i];
                    }
                    if (writeCounts) {
                        w.write(mCounts[i][0] + "\t" + mCounts[i][1] + "\n");
                    }
                }
                if (writeCounts) {
                    w.close();
                }
            }

            return true;
        } catch (Exception e) {
            System.out.println("Failed to load counts: " + e);
            e.printStackTrace();
            return false;
        }
    }

    private long[] fastReadCount(int aArity, long[] aOneGramCounts) {
        if (aArity == 1 || aOneGramCounts == null) {
            return readCount(1);
        } else {
            return new long[]{aOneGramCounts[0] - aOneGramCounts[2] * (aArity - 1), 0, 0};
        }
    }

    private long[] readCount(int aArity) {
        if (aArity <= 0 || aArity > mMaxArity) {
            return null;
        }

        int maxidx;
        long ct = 0, disc_ct = 0, s_ct = 0;

        if (aArity == 1) {
            maxidx = 1;
        } else {
            maxidx = mIndices[aArity].length;
        }

        System.out.println("Beginning count pass over " + maxidx + " " + aArity + "-gram files.  Please wait.");

        for (int i = 0; i < maxidx; i++) {
            BufferedReader r;

            System.out.println("Counting file " + (i + 1) + " / " + maxidx + "...");

            try {
                r = new BufferedReader(new InputStreamReader(new GZIPInputStream(new FileInputStream(generatePath(aArity, i)))));
            } catch (Exception e) {
                System.out.println("Failed to open file: " + e);
                return null;
            }

            String l;

            try {
                while ((l = r.readLine()) != null) {
                    //String[] w = l.split("\t");
                    String[] w = new String[2];
                    w[0] = l.substring(0, l.indexOf("\t"));
                    w[1] = l.substring(l.indexOf("\t") + 1);
                    if (w.length != 2) {
                        return null;
                    }

                    long v = Long.parseLong(w[1]);

                    ct += v;
                    disc_ct++;

                    if (aArity == 1 && w[0].equals("</S>")) {
                        s_ct = v;
                    }

                    //if (ct < old_ct) System.out.println("Long overflow detected...");
                }
            } catch (Exception e) {
            } finally {
                try {
                    r.close();
                } catch (Exception e) {
                }
            }
        }

        System.out.println();
        System.out.println("Counted " + ct + " " + aArity + "-grams (" + disc_ct + " discrete) in " + maxidx + " files.");

        return new long[]{ct, disc_ct, s_ct};
    }

    public long[] getTotalCount(int aArity) {
        if (aArity <= 0 || aArity > mMaxArity) {
            return null;
        }

        return mCounts[aArity].clone();
    }

    public Query executeQuery(Query aQ) {
        if (aQ.getArity() > mMaxArity) {
            return null;
        }
        
        int n = aQ.getQueryString().length(), c;
        
        if ((c = aQ.getQueryString().compareTo(mFilePointers[n].last_query)) < 0) {
            // we've already stepped over this query in the n-gram database.  since we read sequentially,
            // it's obviously not in the database.  return a count of zero.
            aQ.finish(0);
            return aQ;
        } else if (c == 0) {
            // mLastQueryResults holds the cached value...
            aQ.finish(mFilePointers[n].last_query_result);
            return aQ;
        }

        int fn = Math.max(mFilePointers[n].fidx, 0);

        // load the proper file for the queried n-gram.
        // HACK: for 1-grams, skip this step.  there's only one 1-gram file.
        if (n != 1) {
            while (fn < mIndices[n].length && mIndices[n][fn].compareTo(aQ.getQueryString()) < 0) {
                //System.out.println("next file...");
                fn++;
            }

            if (fn == 0) {
                // the query is before the first file in the database.
                // therefore, it is not in the database, and has a count of zero.
                aQ.finish(0);
                return aQ;
            }

            if (openFile(n, --fn) == false) {
                return null;
            }
        }

        // scan the file for the n-gram...
        QueryString ng;

        while (true) {
            String l;
            //System.out.println("entrei");
            try {
                l = mFilePointers[n].fp.readLine();
            } catch (Exception e) {
                System.out.println("Query failed: error while reading: " + e);
                return null;
            }
            if (l == null || l.trim() == "") {
                // we ran off the end of the file; this n-gram isn't in the database.
                aQ.finish(0);
               // System.out.println("terminei");
                return aQ;
            }

            //String[] w = l.split("\t");
            String[] w = new String[2];
            w[0] = l.substring(0, l.indexOf("\t"));
            w[1] = l.substring(l.indexOf("\t") + 1);

            if (w.length != 2) {
                System.out.println("Query failed: invalid format for ngram record " + l);
                return null;
            }

            //ng = new QueryString(w[0].split(" "));
            ng = new QueryString(w[1].split("\t"));
            int r = ng.compareTo(aQ.getQueryString());

            if (r >= 0) {
                // we landed on the n-gram, or skipped over it.  return a count of 0 (if
                // we skipped over it) or the listed count (if we landed on it) and reset
                // last-query to the n-gram we stopped on.
                //long v = Long.parseLong(w[1]);
                long v = Long.parseLong(w[0]);

                aQ.finish((r > 0) ? 0 : v);
                mFilePointers[n].last_query = ng;
                mFilePointers[n].last_query_result = v;
                break;
            }
        }
        return aQ;
    }

    public String getLocation() {
        return mPath;
    }
    private String mPath;
    private QueryString[][] mIndices;
    private long[][] mCounts;
    
    @Override
    public void closeFiles() {
        for(Pointer p : mFilePointers){
            if(p.fp != null){
                try {
                    p.fp.close();
                } catch (IOException ex) {
                    System.out.println("Error: Failed to close the n-gram files.");
                    Logger.getLogger(Gale1TNGramDatabase.class.getName()).log(Level.SEVERE, null, ex);
                }
            }
        }
    }

    private class Pointer {

        public BufferedReader fp = null;
        public int fidx = -1;
        public QueryString last_query = null;
        public long last_query_result = -1;
    }
    private Pointer[] mFilePointers;

    public int getMaximumArity() {
        return mMaxArity;
    }
    private int mMaxArity;

    /* code below this is for testing purposes! */
    public static void main(String[] aArgs) {
        Gale1TNGramDatabase db = null;

        try {
            db = new Gale1TNGramDatabase("/media/web_5gram_v1_1.btw/data", 2);
        } catch (Exception e) {
            System.out.println("Couldn't open ngram database: " + e);
            System.exit(1);
        }

        String[] w = new String[]{"I", "am", "the", "very", "model"};
        String[] v = new String[]{"Jackdaws", "love", "my", "big", "splenectomy"};

        Query q = new Query(w);
        Query r = new Query(v);

        /*db.executeQuery(q);
        db.executeQuery(r);
        System.out.println(q);
        System.out.println(r);*/

        long[] ct = db.getTotalCount(1);
        System.out.println("There are " + ct[0] + " 1-grams in the database, " + ct[1] + " distinct");
        ct = db.getTotalCount(2);
        System.out.println("There are " + ct + " 2-grams in the database, " + ct[1] + " distinct");
    }
}















